package net.cyue.ort.llm.data;

import java.util.*;
import java.util.stream.Collectors;

/**
 * Token列表集合类
 * 提供token序列的管理和分数计算
 */
public class TokenList implements Collection<Token>, Comparable<TokenList>, AutoCloseable {
    private final List<Token> tokenList;
    private double sequenceScore;
    
    public TokenList() {
        this.tokenList = new ArrayList<>();
        this.sequenceScore = 0.0;
    }
    
    public TokenList(List<Token> tokenList) {
        this.tokenList = new ArrayList<>(tokenList);
        this.sequenceScore = tokenList.stream().mapToDouble(Token::getScore).sum();
    }
    
    public TokenList(long[] tokenIDs) {
        this.tokenList = new ArrayList<>();
        for (long tokenID : tokenIDs) {
            this.tokenList.add(new Token(tokenID));
        }
        this.sequenceScore = 0.0;
    }
    
    public TokenList(Long[] tokenIDs) {
        this.tokenList = new ArrayList<>();
        for (long tokenID : tokenIDs) {
            this.tokenList.add(new Token(tokenID));
        }
        this.sequenceScore = 0.0;
    }
    
    /**
     * 获取序列分数
     */
    public double getScore() {
        return sequenceScore;
    }
    
    /**
     * 获取序列分数（别名方法）
     */
    public double getSequenceScore() {
        return sequenceScore;
    }
    
    @Override
    public int compareTo(TokenList other) {
        // 按序列分数降序排序
        return Double.compare(other.sequenceScore, this.sequenceScore);
    }
    
    @Override
    public void close() {
        tokenList.clear();
        sequenceScore = 0.0;
    }
    
    /**
     * 检查是否包含指定token
     */
    public boolean contains(Token token) {
        return tokenList.contains(token);
    }
    
    /**
     * 检查是否包含指定token ID
     */
    public boolean contains(long tokenID) {
        return tokenList.stream().anyMatch(token -> token.getId() == tokenID);
    }
    
    /**
     * 获取指定索引的token
     */
    public Token get(int index) {
        return tokenList.get(index);
    }
    
    /**
     * 获取token（别名方法）
     */
    public Token getToken(int index) {
        return tokenList.get(index);
    }
    
    /**
     * 获取最后一个token
     */
    public Token getLastToken() {
        if (tokenList.isEmpty()) {
            throw new IndexOutOfBoundsException("TokenList is empty");
        }
        return tokenList.get(tokenList.size() - 1);
    }
    
    /**
     * 获取所有token ID数组
     */
    public long[] getTokenIDs() {
        return tokenList.stream().mapToLong(Token::getId).toArray();
    }
    
    /**
     * 获取所有token ID列表
     */
    public List<Long> getTokenIDList() {
        return tokenList.stream().map(Token::getId).collect(Collectors.toList());
    }
    
    /**
     * 添加token
     */
    @Override
    public boolean add(Token token) {
        tokenList.add(token);
        sequenceScore += token.getScore();
        return true;
    }
    
    /**
     * 添加token（别名方法）
     */
    public void addToken(Token token) {
        add(token);
    }
    
    /**
     * 添加另一个TokenList中的所有token
     */
    public void addAll(TokenList tokenList) {
        this.tokenList.addAll(tokenList.tokenList);
        this.sequenceScore += tokenList.sequenceScore;
    }
    
    /**
     * 添加List中的所有token
     */
    public void addAll(List<Token> tokenList) {
        for (Token token : tokenList) {
            add(token);
        }
    }
    
    /**
     * 添加TokenSet中的所有token
     */
    public void addAll(TokenSet tokenSet) {
        for (Token token : tokenSet) {
            add(token);
        }
    }
    
    /**
     * 添加Set中的所有token
     */
    public void addAll(Set<Token> tokenSet) {
        for (Token token : tokenSet) {
            add(token);
        }
    }
    
    /**
     * 添加TokenList（别名方法）
     */
    public void addTokenList(TokenList tokenList) {
        addAll(tokenList);
    }
    
    /**
     * 添加List（别名方法）
     */
    public void addTokenList(List<Token> tokenList) {
        addAll(tokenList);
    }
    
    /**
     * 添加TokenSet（别名方法）
     */
    public void addTokenSet(TokenSet tokenSet) {
        addAll(tokenSet);
    }
    
    /**
     * 添加Set（别名方法）
     */
    public void addTokenSet(Set<Token> tokenSet) {
        addAll(tokenSet);
    }
    
    /**
     * 排序token列表（按分数降序）
     */
    public void sort() {
        Collections.sort(tokenList);
    }
    
    // Collection接口实现
    
    @Override
    public int size() {
        return tokenList.size();
    }
    
    @Override
    public boolean isEmpty() {
        return tokenList.isEmpty();
    }
    
    @Override
    public boolean contains(Object o) {
        return tokenList.contains(o);
    }
    
    @Override
    public Iterator<Token> iterator() {
        return tokenList.iterator();
    }
    
    @Override
    public Object[] toArray() {
        return tokenList.toArray();
    }
    
    @Override
    public <T> T[] toArray(T[] a) {
        return tokenList.toArray(a);
    }
    
    @Override
    public boolean remove(Object o) {
        if (o instanceof Token) {
            Token token = (Token) o;
            boolean removed = tokenList.remove(token);
            if (removed) {
                sequenceScore -= token.getScore();
            }
            return removed;
        }
        return false;
    }
    
    @Override
    public boolean containsAll(Collection<?> c) {
        return tokenList.containsAll(c);
    }
    
    @Override
    public boolean addAll(Collection<? extends Token> c) {
        boolean modified = false;
        for (Token token : c) {
            if (add(token)) {
                modified = true;
            }
        }
        return modified;
    }
    
    @Override
    public boolean removeAll(Collection<?> c) {
        boolean modified = false;
        for (Object o : c) {
            if (remove(o)) {
                modified = true;
            }
        }
        return modified;
    }
    
    @Override
    public boolean retainAll(Collection<?> c) {
        boolean modified = false;
        Iterator<Token> it = tokenList.iterator();
        while (it.hasNext()) {
            Token token = it.next();
            if (!c.contains(token)) {
                it.remove();
                sequenceScore -= token.getScore();
                modified = true;
            }
        }
        return modified;
    }
    
    @Override
    public void clear() {
        tokenList.clear();
        sequenceScore = 0.0;
    }
}

